# compiler/ast.py
from _ast import Expression
from enum import Enum

from src.compiler.Lexer.tokens import TokenType
from src.compiler.Parser.grammar import Terminal


class Node:
    """AST节点基类"""

    def __repr__(self):
        # 使用默认的对象表示，避免递归
        return f"<{self.__class__.__name__} at {hex(id(self))}>"

    def __str__(self):
        # 提供基本的字符串表示
        return self.__repr__()

class Literal(Expression):
    """字面量，如数字、字符串"""
    def __init__(self, value, terminal_type):
        self.value = value
        self.terminal_type = terminal_type

    def __str__(self):
        if self.terminal_type == Terminal.STRING:
            return f"'{self.value}'"
        return str(self.value)

class Assignment(Node):
    """UPDATE语句中的赋值表达式，如 column = value"""
    def __init__(self, column, value):
        self.column = column  # Identifier 对象
        self.value = value    # Expression 对象

    def __str__(self):
        return f"{self.column} = {self.value}"
class Statement(Node):
    """语句基类"""
    pass


class SelectStatement(Statement):
    """SELECT语句"""

    def __init__(self, target_list, from_table, where_clause=None):
        self.target_list = target_list  # 目标列列表，如['id', 'name'] 或 ['*']
        self.from_table = from_table  # 表名，如'students' (Identifier对象)
        self.where_clause = where_clause  # WHERE条件表达式 (Expression对象)

    def __str__(self):
        where_str = f" WHERE {self.where_clause}" if self.where_clause else ""
        targets_str = ", ".join([str(t) for t in self.target_list])
        return f"SELECT {targets_str} FROM {self.from_table}{where_str}"


class CreateTableStatement(Statement):
    """CREATE TABLE语句"""

    def __init__(self, table_name, columns):
        self.table_name = table_name  # 表名 (Identifier对象)
        self.columns = columns  # list of (column_name, data_type) 元组

    def __str__(self):
        columns_str = ", ".join([f"{col[0]} {col[1]}" for col in self.columns])
        return f"CREATE TABLE {self.table_name} ({columns_str})"


class InsertStatement(Statement):
    """INSERT语句"""

    def __init__(self, table_name, column_names, values):
        self.table_name = table_name  # 表名 (Identifier对象)
        self.column_names = column_names  # 列名列表 [Identifier]
        self.values = values  # 值列表 [Expression]

    def __str__(self):
        columns_str = ", ".join([str(col) for col in self.column_names])
        values_str = ", ".join([str(val) for val in self.values])
        return f"INSERT INTO {self.table_name} ({columns_str}) VALUES ({values_str})"


class DeleteStatement(Statement):
    """DELETE语句"""

    def __init__(self, table_name, where_clause=None):
        self.table_name = table_name  # 表名 (Identifier对象)
        self.where_clause = where_clause  # WHERE条件表达式 (Expression对象)

    def __str__(self):
        where_str = f" WHERE {self.where_clause}" if self.where_clause else ""
        return f"DELETE FROM {self.table_name}{where_str}"


class UpdateStatement(Statement):
    """UPDATE语句（为未来扩展准备）"""

    def __init__(self, table_name, set_list, where_clause=None):
        self.table_name = table_name  # 表名
        self.set_list = set_list  # SET子句列表 [(column, value)]
        self.where_clause = where_clause  # WHERE条件表达式

    def __str__(self):
        set_str = ", ".join([f"{col} = {val}" for col, val in self.set_list])
        where_str = f" WHERE {self.where_clause}" if self.where_clause else ""
        return f"UPDATE {self.table_name} SET {set_str}{where_str}"

class WhereClause(Node):
    """WHERE 子句"""
    def __init__(self, condition):
        self.condition = condition # Expression 对象

    def __str__(self):
        return str(self.condition)

class Expression(Node):
    """表达式基类"""
    pass


class BinaryExpression(Expression):
    """二元表达式，如 age > 18"""

    def __init__(self,left, op, right):
        self.op = op  # 操作符，如 '>', '=', '+', '-'
        self.left = left  # 左操作数 (Expression)
        self.right = right  # 右操作数 (Expression)

    def __str__(self):
        # 确保操作符能正确显示其字符串值
        op_str = self.op.value if isinstance(self.op, Terminal) else str(self.op)
        return f"({self.left} {op_str} {self.right})"



class LogicalExpression(Expression):
    """逻辑表达式，如 age > 18 AND name = 'John'"""

    def __init__(self, op, left, right):
        self.op = op  # 逻辑操作符，如 'AND', 'OR'
        self.left = left  # 左表达式
        self.right = right  # 右表达式

    def __str__(self):
        return f"({self.left} {self.op} {self.right})"


class UnaryExpression(Expression):
    """一元表达式，如 -price 或 NOT active"""

    def __init__(self, op, operand):
        self.op = op  # 操作符，如 '-', 'NOT'
        self.operand = operand  # 操作数

    def __str__(self):
        op_str = self.op.value if isinstance(self.op, Terminal) else str(self.op)
        return f"({op_str}{self.operand})"


class FunctionCall(Expression):
    """函数调用表达式，如 COUNT(*) 或 MAX(price)"""

    def __init__(self, function_name, arguments):
        self.function_name = function_name  # 函数名
        self.arguments = arguments  # 参数列表

    def __str__(self):
        args_str = ", ".join([str(arg) for arg in self.arguments])
        return f"{self.function_name}({args_str})"


class Identifier(Expression):
    """标识符，如表名、列名"""
    def __init__(self, name):
        self.name = name

    def __str__(self):
        return str(self.name)


class Constant(Expression):
    """常量/字面量，如数字、字符串"""
    def __init__(self, value):
        self.value = value

    def __str__(self):
        # 如果是字符串，加上引号
        if isinstance(self.value, str):
            return f"'{self.value}'"
        return str(self.value)


class ColumnRef(Expression):
    """列引用，可能包含表名，如 table.column"""

    def __init__(self, table_name, column_name):
        self.table_name = table_name  # 表名（可选）
        self.column_name = column_name  # 列名

    def __str__(self):
        if self.table_name:
            return f"{self.table_name}.{self.column_name}"
        return self.column_name


class TypeCast(Expression):
    """类型转换表达式，如 price::INT"""

    def __init__(self, expression, target_type):
        self.expression = expression  # 要转换的表达式
        self.target_type = target_type  # 目标类型

    def __str__(self):
        return f"{self.expression}::{self.target_type}"


class BetweenExpression(Expression):
    """BETWEEN表达式，如 age BETWEEN 18 AND 65"""

    def __init__(self, expression, lower_bound, upper_bound):
        self.expression = expression  # 要检查的表达式
        self.lower_bound = lower_bound  # 下界
        self.upper_bound = upper_bound  # 上界

    def __str__(self):
        return f"{self.expression} BETWEEN {self.lower_bound} AND {self.upper_bound}"


class InExpression(Expression):
    """IN表达式，如 id IN (1, 2, 3)"""

    def __init__(self, expression, values):
        self.expression = expression  # 要检查的表达式
        self.values = values  # 值列表

    def __str__(self):
        values_str = ", ".join([str(val) for val in self.values])
        return f"{self.expression} IN ({values_str})"


class IsNullExpression(Expression):
    """IS NULL表达式，如 name IS NULL"""

    def __init__(self, expression, is_null=True):
        self.expression = expression  # 要检查的表达式
        self.is_null = is_null  # 是否为NULL检查

    def __str__(self):
        if self.is_null:
            return f"{self.expression} IS NULL"
        else:
            return f"{self.expression} IS NOT NULL"


class CaseExpression(Expression):
    """CASE表达式"""

    def __init__(self, cases, else_result=None):
        self.cases = cases  # [(condition, result)] 列表
        self.else_result = else_result  # ELSE结果（可选）

    def __str__(self):
        case_str = "CASE "
        for condition, result in self.cases:
            case_str += f"WHEN {condition} THEN {result} "
        if self.else_result:
            case_str += f"ELSE {self.else_result} "
        case_str += "END"
        return case_str


# 数据类型定义
class DataType(Node):
    """数据类型基类"""

    def __init__(self, name, length=None):
        self.name = name
        self.length = length

    def __str__(self):
        if self.length:
            return f"{self.name}({self.length})"
        return self.name


class IntType(DataType):
    """INTEGER类型"""

    def __init__(self):
        super().__init__("INTEGER")


class VarcharType(DataType):
    """VARCHAR类型"""

    def __init__(self, length=255):
        super().__init__("VARCHAR", length)


class FloatType(DataType):
    """FLOAT类型"""

    def __init__(self):
        super().__init__("FLOAT")


class BooleanType(DataType):
    """BOOLEAN类型"""

    def __init__(self):
        super().__init__("BOOLEAN")


class DateType(DataType):
    """DATE类型"""

    def __init__(self):
        super().__init__("DATE")


# 表约束定义
class Constraint(Node):
    """约束基类"""
    pass


class PrimaryKeyConstraint(Constraint):
    """主键约束"""

    def __init__(self, columns):
        self.columns = columns  # 列名列表

    def __str__(self):
        cols_str = ", ".join(self.columns)
        return f"PRIMARY KEY ({cols_str})"


class ForeignKeyConstraint(Constraint):
    """外键约束"""

    def __init__(self, columns, ref_table, ref_columns):
        self.columns = columns  # 本表列名
        self.ref_table = ref_table  # 引用表名
        self.ref_columns = ref_columns  # 引用列名

    def __str__(self):
        cols_str = ", ".join(self.columns)
        ref_cols_str = ", ".join(self.ref_columns)
        return f"FOREIGN KEY ({cols_str}) REFERENCES {self.ref_table}({ref_cols_str})"


class UniqueConstraint(Constraint):
    """唯一约束"""

    def __init__(self, columns):
        self.columns = columns

    def __str__(self):
        cols_str = ", ".join(self.columns)
        return f"UNIQUE ({cols_str})"


class NotNullConstraint(Constraint):
    """非空约束"""

    def __init__(self, column):
        self.column = column

    def __str__(self):
        return f"NOT NULL"


# 表定义（包含约束）
class TableDefinition(Node):
    """表定义"""

    def __init__(self, name, columns, constraints=None):
        self.name = name
        self.columns = columns  # [(name, data_type)] 列表
        self.constraints = constraints or []  # 约束列表

    def __str__(self):
        elements = []
        for col_name, data_type in self.columns:
            elements.append(f"{col_name} {data_type}")
        elements.extend([str(constraint) for constraint in self.constraints])
        return f"CREATE TABLE {self.name} (\n    " + ",\n    ".join(elements) + "\n)"